#!/user/bin/env python3
# -*- coding: utf-8 -*-
"""
--------------------------------------
    Author:     JiChao_Song
    Date  :     2022/2/14 16:12
    Desc  :
--------------------------------------
"""

import datetime
import os
import platform
import subprocess
import sys
import time
import uuid

# from extends.flask_locust.main import main as locusts_main
from django.conf.urls import url
from django.core.cache import cache
from django.urls import path
from django_http_log.common.request import BaseAPIView
from django_http_log.common.response import SuccessResponse, SystemExceptionResponse, CommonResponse, \
    ParamsErrorResponse

from common.enums.response_enums import CommonResultEnums
from common.utils.log_utils import logger
# from extends.flask_locust.requests import stat_requests
from extends.flask_locust.requests import stat_requests, stat_stop
from platform_auth.settings import APIDOC_DECORATOR


class LocustStatRequests(BaseAPIView):

    def get(self, request, *args, **kwargs):

        try:
            system_info = platform.system()
            check = ''

            if system_info == 'Windows':
                check = os.popen('tasklist | findstr "locust"')
            if system_info == 'Linux':
                check = os.popen('ps -ef | grep locust')

            if check.read() != '':
                data: dict = stat_requests()
                return SuccessResponse(result = data)
            return CommonResponse(codeEnums = CommonResultEnums.LOCUSTS_PROCESS_NOT_STARTING)

        except Exception as e:

            logger.error(e)

            return SystemExceptionResponse()


class LocustStatStop(BaseAPIView):

    def get(self, request, *args, **kwargs):

        try:
            system_info = platform.system()
            check = ''

            if system_info == 'Windows':
                check = os.popen('tasklist | findstr "locust"')
            if system_info == 'Linux':
                check = os.popen('ps -ef | grep locust')

            if check.read() != '':
                data: dict = stat_stop()
                return SuccessResponse(result = data)
            return CommonResponse(codeEnums = CommonResultEnums.LOCUSTS_PROCESS_NOT_STARTING)

        except Exception as e:

            logger.error(e)

            return SystemExceptionResponse()


class LocustStatCharts(BaseAPIView):

    def get(self, request, *args, **kwargs):

        try:
            stats_history = {}
            system_info = platform.system()
            check = ''

            if system_info == 'Windows':
                check = os.popen('tasklist | findstr "locust"')
            if system_info == 'Linux':
                check = os.popen('ps -ef | grep locust')

            if check.read() != '':
                data: dict = stat_requests()
                time = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
                if data.get('state') and data.get('state') == 'stopped':
                    stats_history['time'] = time
                    stats_history['user_count'] = {'value': None}
                    stats_history['current_rps'] = {'value': None, 'users': 0}
                    stats_history['current_fail_per_sec'] = {'value': None, 'users': 0}
                    stats_history['response_time_percentile_50'] = {'value': None, 'users': 0}
                    stats_history['response_time_percentile_95'] = {'value': None, 'users': 0}
                    return SuccessResponse(result = stats_history)
                stats_history['markers'] = {'xAxis': time}
                total = data.get('stats')[len(data.get('stats')) - 1 if len(data.get('stats')) >= 1 else 0]

                if total and total.get('num_requests') < 1:
                    return SuccessResponse(result = stats_history)
                stats_history['time'] = time
                stats_history['user_count'] = {'value': data.get('user_count')}
                stats_history['current_rps'] = {'value': total.get('current_rps'), 'users': data.get('user_count')}
                stats_history['current_fail_per_sec'] = {'value': total.get('current_fail_per_sec'),
                                                         'users': data.get('user_count')}
                stats_history['response_time_percentile_50'] = {'value': data.get('response_time_percentile_50'),
                                                                'users': data.get('user_count')}
                stats_history['response_time_percentile_95'] = {'value': data.get('response_time_percentile_95'),
                                                                'users': data.get('user_count')}

                return SuccessResponse(result = stats_history)
            return SuccessResponse()
        except Exception as e:

            logger.error(e)

            return SystemExceptionResponse()


class LocustStatsStartAPIView(BaseAPIView):

    @APIDOC_DECORATOR.checkParams
    @APIDOC_DECORATOR.start
    @APIDOC_DECORATOR.ApiTag(value = 'locust', desc = '性能测试')
    @APIDOC_DECORATOR.ApiOperation(method = 'POST', path = '/locust/stats/start', summary = '启动性能测试')
    @APIDOC_DECORATOR.ApiParams(value = 'host', require = True, type = 'string', desc = '主机')
    @APIDOC_DECORATOR.ApiParams(value = 'userCount', require = True, type = 'int', desc = '模拟用户数')
    @APIDOC_DECORATOR.ApiParams(value = 'spawnUser', type = 'int', require = True, desc = '每秒用户数')
    def post(self, request, *args, **kwargs):
        try:
            body: dict = request.data.copy()
            host = body.get('host')
            userCount = body.get('userCount')
            spawnUser = body.get('spawnUser')
            system_info = platform.system()
            check = ''

            if system_info == 'Windows':
                check = os.popen('tasklist | findstr "locust"')
            if system_info == 'Linux':
                check = os.popen('ps -ef | grep locust')

            if check.read() != '':
                return CommonResponse(codeEnums = CommonResultEnums.LOCUSTS_PROCESS_STARTING)
            locust_cmd_str = 'locust -f ./extends/flask_locust/user.py --web-host="0.0.0.0" --autostart'
            if userCount:
                locust_cmd_str += f' --users={userCount}'
            if spawnUser:
                locust_cmd_str += f' --spawn-rate={spawnUser}'

            if host:
                locust_cmd_str += f' --host="{host}"'

            logger.info(f'性能测试开始: {locust_cmd_str}')

            subprocess.Popen(locust_cmd_str)

            return SuccessResponse()

        except Exception as e:

            logger.error(e)

            return SystemExceptionResponse()

    def get(self, request, *args, **kwargs):
        try:
            system_info = platform.system()
            check = ''

            if system_info == 'Windows':
                check = os.popen('tasklist | findstr "locust"')
            if system_info == 'Linux':
                check = os.popen('ps -ef | grep locust')

            if check.read() != '':
                return CommonResponse(codeEnums = CommonResultEnums.LOCUSTS_PROCESS_STARTING)
            locust_cmd_str = 'locust -f ./extends/flask_locust/user.py --web-host="0.0.0.0" --autostart'

            subprocess.Popen(locust_cmd_str)

            return SuccessResponse()

        except Exception as e:

            logger.error(e)

            return SystemExceptionResponse()


class LocustStatsCloseAPIView(BaseAPIView):

    def get(self, request, *args, **kwargs):
        try:

            system_info = platform.system()

            if system_info == 'Windows':

                os.system(f'taskkill -f /IM locust.exe')

            elif system_info == 'Linux':

                os.system("ps -ef | grep locust|grep -v 'grep ' | awk '{print $2}'|xargs kill -9")

            return SuccessResponse()

        except Exception as e:

            logger.error(e)

            return SystemExceptionResponse()


urlpatterns = [
    url('stats/requests', LocustStatRequests.as_view()),
    url('stats/stop', LocustStatStop.as_view()),
    url('stats/charts', LocustStatCharts.as_view()),
    url('stats/start', LocustStatsStartAPIView.as_view()),
    url('stats/close', LocustStatsCloseAPIView.as_view()),
]
